library(tidyverse)
library(readxl)
library(rdrobust)
library(gridExtra)
library(grid)

# ---- load dataset ------
load("data/shijizhengfa.RData")
load("data/shengjizhengfa.RData")

# Robustness check using posts by city-level Political & Legal Affairs Committees (政法委)
# Mobilizing the Grid during the "Dynamic Zero COVID" period

test_list <- read_xlsx("data/sporadic_test.xlsx") |>
  mutate(First_Infected = as.Date(First_Infected))

# Subset used for local (city-level) tests only (exclude Yili)
test_list2 <- test_list |> 
  filter(City != "Yili")

process_data_local_zhengfawei <- function(test_list2, shijizhengfa) {
  all_results <- data.frame() 
  
  for (i in 1:nrow(test_list2)) {
    First_Infected <- as.Date(test_list2$First_Infected[i])
    
    City <- test_list2$City[i]
    Scale <- test_list2$Scale[i]
    Regional_spreading <- test_list2$Regional_spreading[i]
    National_spreading <- test_list2$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date <- First_Infected %m+% months(12)
    
    City_filter <- test_list2$City[i]
    
    rd_data <- shijizhengfa |>
      filter(City %in% City_filter) |>
      filter(date >= start_date & date <= end_date)
    
    rd_data <- rd_data |>
      mutate(difference = as.integer(date - First_Infected),
             others = total_count - daily_grid)
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    rd_control <- rdrobust(rd_data$daily_grid,
                           rd_data$difference,
                           covs = rd_data$others,
                           p = 1,
                           kernel = "triangular",
                           bwselect = "mserd")
    
    conventional_coef <- rd_control[["coef"]]["Conventional", "Coeff"] / daily_total_mean
    conventional_ci_lower <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
    conventional_ci_upper <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    
    rd_results <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = conventional_coef,
      CI_Lower = conventional_ci_lower,
      CI_Upper = conventional_ci_upper
    )
    
    all_results <- bind_rows(all_results, rd_results)
  }
  
  return(all_results)
}

local_results_daily_zhengfawei <- process_data_local_zhengfawei(test_list2, shijizhengfa)

print(local_results_daily_zhengfawei)

save(local_results_daily_zhengfawei, file = "generated_data/dynamic_RDD_results/local_results_daily_zhengfawei.RData")

# Regional mobilization, daily, no lag
process_data_regional_zhengfawei <- function(test_list, shengjizhengfa) {
  all_results <- data.frame() 
  
  for (i in 1:nrow(test_list)) {
    First_Infected <- as.Date(test_list$First_Infected[i])
    
    City <- test_list$City[i]
    Scale <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date <- First_Infected %m+% months(12)
    
    peripheral_filter <- unlist(strsplit(test_list$Peripheral[i], ","))
    
    rd_data <- shengjizhengfa |>
      filter(province %in% peripheral_filter) |>
      filter(date >= start_date & date <= end_date)
    
    rd_data$difference <- as.integer(rd_data$date - First_Infected)
    rd_data$week <- NA
    
    rd_data$others <- rd_data$total_count - rd_data$daily_grid
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    rd_control <- rdrobust(rd_data$daily_grid,
                           rd_data$difference,
                           covs = rd_data$others,
                           p = 1,
                           kernel = "triangular",
                           bwselect = "mserd",
                           c = 0)
    
    conventional_coef <- rd_control[["coef"]]["Conventional", "Coeff"] / daily_total_mean
    conventional_ci_lower <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
    conventional_ci_upper <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    
    rd_results <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = conventional_coef,
      CI_Lower = conventional_ci_lower,
      CI_Upper = conventional_ci_upper
    )
    
    all_results <- bind_rows(all_results, rd_results)
  }
  
  return(all_results)
}

regional_results_daily_zhengfawei <- process_data_regional_zhengfawei(test_list, shengjizhengfa)

print(regional_results_daily_zhengfawei)

save(regional_results_daily_zhengfawei, file = "generated_data/dynamic_RDD_results/regional_results_daily_zhengfawei.RData")

# National mobilization, daily, no lag
process_data_national_zhengfawei <- function(test_list, shengjizhengfa) {
  all_results <- vector("list", nrow(test_list))
  
  for (i in seq_len(nrow(test_list))) {
    First_Infected <- as.Date(test_list$First_Infected[i])
    City <- test_list$City[i]
    Scale <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date   <- First_Infected %m+% months(12)
    
    national_raw <- test_list$National_level[i]
    national <- if (is.na(national_raw) || national_raw == "") character(0) else {
      trimws(unlist(strsplit(national_raw, ",")))
    }
    
    rd_data <- shengjizhengfa |>
      mutate(date = as.Date(date)) |>
      filter(!(province %in% national)) |>
      filter(date >= start_date & date <= end_date) |>
      group_by(date) |>
      summarise(
        total_count = sum(total_count, na.rm = TRUE),
        daily_grid  = sum(daily_grid,  na.rm = TRUE),
        .groups = "drop"
      ) |>
      complete(date = seq.Date(start_date, end_date, by = "day"),
               fill = list(total_count = 0L, daily_grid = 0L)) |>
      arrange(date) |>
      mutate(
        difference = as.integer(date - First_Infected),
        others     = pmax(total_count - daily_grid, 0L) 
      )
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    if (is.na(daily_total_mean) || daily_total_mean == 0) {
      coef <- NA_real_; ci_l <- NA_real_; ci_u <- NA_real_
    } else {
      rd_control <- rdrobust(
        rd_data$daily_grid,
        rd_data$difference,
        covs = rd_data$others,
        p = 1,
        kernel = "triangular",
        bwselect = "mserd",
        c = 0
      )
      coef <- rd_control[["coef"]]["Conventional", "Coeff"]      / daily_total_mean
      ci_l <- rd_control[["ci"]]["Conventional", "CI Lower"]     / daily_total_mean
      ci_u <- rd_control[["ci"]]["Conventional", "CI Upper"]     / daily_total_mean
    }
    
    all_results[[i]] <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = coef,
      CI_Lower = ci_l,
      CI_Upper = ci_u,
      stringsAsFactors = FALSE
    )
  }
  
  dplyr::bind_rows(all_results)
}

national_results_daily_zhengfawei <- process_data_national_zhengfawei(test_list, shengjizhengfa)

print(national_results_daily_zhengfawei)

save(national_results_daily_zhengfawei, file = "generated_data/dynamic_RDD_results/national_results_daily_zhengfawei.RData")
# Add a row for Yili (kept as NA because local test_list2 excluded Yili earlier)
Yili_row <- data.frame(
  City = "Yili",
  Scale = 4,
  Regional_spreading = 1,
  National_spreading = 0,
  First_Infected = "2022-07-30",
  Coefficient = NA,
  CI_Lower = NA,
  CI_Upper = NA
)

beihai_index <- which(local_results_daily_zhengfawei$City == "Beihai")

local_results_daily_zhengfawei <- rbind(
  local_results_daily_zhengfawei[1:beihai_index, ],
  Yili_row,
  local_results_daily_zhengfawei[(beihai_index + 1):nrow(local_results_daily_zhengfawei), ]
)

local_results_daily_zhengfawei

# Plot local mobilization, daily, no lag
local_results_daily_zhengfawei <- local_results_daily_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

local_results_daily_zhengfawei$First_Infected <- format(as.Date(local_results_daily_zhengfawei$First_Infected), "%d %b, %Y")

local_results_daily_zhengfawei <- local_results_daily_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

local_mobilization_zhengfawei <- ggplot(local_results_daily_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  
  scale_x_continuous(breaks = local_results_daily_zhengfawei$Order, labels = local_results_daily_zhengfawei$First_Infected) +
  labs(title = "",
       x = "",
       y = "Local Mobilization") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        # x-axis text hidden for compact layout across facets
        axis.text.y = element_text(size = 16),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

print(local_mobilization_zhengfawei)

# Plot regional mobilization, daily
regional_results_daily_zhengfawei <- regional_results_daily_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

regional_results_daily_zhengfawei$First_Infected <- format(as.Date(regional_results_daily_zhengfawei$First_Infected), "%d %b, %Y")

regional_results_daily_zhengfawei <- regional_results_daily_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

regional_mobilization_zhengfawei <- ggplot(regional_results_daily_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  
  scale_x_continuous(breaks = regional_results_daily_zhengfawei$Order, labels = regional_results_daily_zhengfawei$City) +
  
  labs(title = "",
       x = "",
       y = "Regional Mobilization") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        # Hide axis text for a cleaner middle panel
        axis.text.y = element_blank(),
        axis.ticks.y = element_blank(),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

print(regional_mobilization_zhengfawei)

# Plot national mobilization, daily (no lag)
national_results_daily_zhengfawei <- national_results_daily_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

national_results_daily_zhengfawei$First_Infected <- format(as.Date(national_results_daily_zhengfawei$First_Infected), "%d %b, %Y")

national_results_daily_zhengfawei <- national_results_daily_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

national_mobilization_zhengfawei <- ggplot(national_results_daily_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  # Empty vs solid points by spreading
  scale_x_continuous(breaks = national_results_daily_zhengfawei$Order, , labels = national_results_daily_zhengfawei$City, position = "top") +
  labs(title = "",
       x = "",
       y = "National Mobilization") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        axis.text.y = element_text(size = 16),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

national_mobilization_zhengfawei


title_textA11_1 <- "FIGURE A10-1. Instant Mobilization Effects at Local, Regional, and National Levels (Robustness Checks)"
title_grobA11_1 <- grobTree(
  textGrob(title_textA11_1, 
           gp = gpar(fontsize = 24, fontface = "bold"), y = -0.6)
)

# Figure notes (bottom)
note_textA11_1 <- "\nThis plot presents the instant mobilization effects at all three levels. The size of each point indicates the scale of the local outbreak. Solid points \nrepresent outbreaks that spread to other cities, while hollow points represent those that did not. Points colored in black indicate the statistical \nsignificance of each mobilization effect. The left side indicates the date the first infection was detected, and the right side indicates the city \nwhere the outbreak occurred."
note_grobA11_1 <- textGrob(note_textA11_1, 
                           gp = gpar(fontsize = 16), 
                           x = 0.1, hjust = 0)


figureA11_1 <- arrangeGrob(
  local_mobilization_zhengfawei,
  regional_mobilization_zhengfawei,
  national_mobilization_zhengfawei,
  ncol = 3,
  #  top = title_grob5,
  #  bottom = note_grob5,
  widths = c(1.08, 0.84, 1.08) 
)

figureA11_1
ggsave(filename = "plots/FIGUREA11_1.png", plot = figureA11_1, width = 16.5, height = 9.2, units = "in", dpi = 200)


# Regional mobilization with lags
# Lag = 7 days
process_data_regional_lag7_zhengfawei<- function(test_list, shengjizhengfa) {
  all_results <- data.frame() 
  
  for (i in 1:nrow(test_list)) {
    First_Infected <- as.Date(test_list$First_Infected[i])
    
    City <- test_list$City[i]
    Scale <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date <- First_Infected %m+% months(12)
    
    peripheral_filter <- unlist(strsplit(test_list$Peripheral[i], ","))
    
    rd_data <- shengjizhengfa |>
      filter(province %in% peripheral_filter) |>
      filter(date >= start_date & date <= end_date)
    
    rd_data$difference <- as.integer(rd_data$date - First_Infected)
    rd_data$week <- NA
    
    rd_data$others <- rd_data$total_count - rd_data$daily_grid
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    rd_control <- rdrobust(rd_data$daily_grid,
                           rd_data$difference,
                           covs = rd_data$others,
                           p = 1,
                           kernel = "triangular",
                           bwselect = "mserd",
                           c = 7)
    
    conventional_coef <- rd_control[["coef"]]["Conventional", "Coeff"] / daily_total_mean
    conventional_ci_lower <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
    conventional_ci_upper <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    
    rd_results <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = conventional_coef,
      CI_Lower = conventional_ci_lower,
      CI_Upper = conventional_ci_upper
    )
    
    all_results <- bind_rows(all_results, rd_results)
  }
  
  return(all_results)
}

regional_results_daily_lag7_zhengfawei <- process_data_regional_lag7_zhengfawei(test_list, shengjizhengfa)

print(regional_results_daily_lag7_zhengfawei)

save(regional_results_daily_lag7_zhengfawei, file = "generated_data/dynamic_RDD_results/regional_results_daily_lag7_zhengfawei.RData")

# Lag = 14 days
process_data_regional_lag14_zhengfawei<- function(test_list, shengjizhengfa) {
  all_results <- data.frame() 
  
  for (i in 1:nrow(test_list)) {
    First_Infected <- as.Date(test_list$First_Infected[i])
    
    City <- test_list$City[i]
    Scale <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date <- First_Infected %m+% months(12)
    
    peripheral_filter <- unlist(strsplit(test_list$Peripheral[i], ","))
    
    rd_data <- shengjizhengfa |>
      filter(province %in% peripheral_filter) |>
      filter(date >= start_date & date <= end_date)
    
    rd_data$difference <- as.integer(rd_data$date - First_Infected)
    rd_data$week <- NA
    
    rd_data$others <- rd_data$total_count - rd_data$daily_grid
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    rd_control <- rdrobust(rd_data$daily_grid,
                           rd_data$difference,
                           covs = rd_data$others,
                           p = 1,
                           kernel = "triangular",
                           bwselect = "mserd",
                           c = 14)
    
    conventional_coef <- rd_control[["coef"]]["Conventional", "Coeff"] / daily_total_mean
    conventional_ci_lower <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
    conventional_ci_upper <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    
    rd_results <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = conventional_coef,
      CI_Lower = conventional_ci_lower,
      CI_Upper = conventional_ci_upper
    )
    
    all_results <- bind_rows(all_results, rd_results)
  }
  
  return(all_results)
}

regional_results_daily_lag14_zhengfawei <- process_data_regional_lag14_zhengfawei(test_list, shengjizhengfa)

print(regional_results_daily_lag14_zhengfawei)

save(regional_results_daily_lag14_zhengfawei, file = "generated_data/dynamic_RDD_results/regional_results_daily_lag14_zhengfawei.RData")


# Lag = 21 days
process_data_regional_lag21_zhengfawei<- function(test_list, shengjizhengfa) {
  all_results <- data.frame() 
  
  for (i in 1:nrow(test_list)) {
    First_Infected <- as.Date(test_list$First_Infected[i])
    
    City <- test_list$City[i]
    Scale <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date <- First_Infected %m+% months(12)
    
    peripheral_filter <- unlist(strsplit(test_list$Peripheral[i], ","))
    
    rd_data <- shengjizhengfa |>
      filter(province %in% peripheral_filter) |>
      filter(date >= start_date & date <= end_date)
    
    rd_data$difference <- as.integer(rd_data$date - First_Infected)
    rd_data$week <- NA
    
    rd_data$others <- rd_data$total_count - rd_data$daily_grid
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    rd_control <- rdrobust(rd_data$daily_grid,
                           rd_data$difference,
                           covs = rd_data$others,
                           p = 1,
                           kernel = "triangular",
                           bwselect = "mserd",
                           c = 21)
    
    conventional_coef <- rd_control[["coef"]]["Conventional", "Coeff"] / daily_total_mean
    conventional_ci_lower <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
    conventional_ci_upper <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    
    rd_results <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = conventional_coef,
      CI_Lower = conventional_ci_lower,
      CI_Upper = conventional_ci_upper
    )
    
    all_results <- bind_rows(all_results, rd_results)
  }
  
  return(all_results)
}

regional_results_daily_lag21_zhengfawei <- process_data_regional_lag21_zhengfawei(test_list, shengjizhengfa)

print(regional_results_daily_lag21_zhengfawei)

save(regional_results_daily_lag21_zhengfawei, file = "generated_data/dynamic_RDD_results/regional_results_daily_lag21_zhengfawei.RData")

# Plot regional mobilization, daily, lag 7 days
regional_results_daily_lag7_zhengfawei <- regional_results_daily_lag7_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

regional_results_daily_lag7_zhengfawei$First_Infected <- format(as.Date(regional_results_daily_lag7_zhengfawei$First_Infected), "%d %b, %Y")

regional_results_daily_lag7_zhengfawei <- regional_results_daily_lag7_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

regional_mobilization_lag7_zhengfawei <- ggplot(regional_results_daily_lag7_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  
  scale_x_continuous(breaks = regional_results_daily_lag7_zhengfawei$Order, labels = regional_results_daily_lag7_zhengfawei$First_Infected) +
  labs(title = "",
       x = "",
       y = "Regional Mobilization, Lagged 7 Days") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        axis.text.y = element_text(size = 16),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

print(regional_mobilization_lag7_zhengfawei)

# Plot regional mobilization, daily, lagged 14 days
regional_results_daily_lag14_zhengfawei <- regional_results_daily_lag14_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

regional_results_daily_lag14_zhengfawei$First_Infected <- format(as.Date(regional_results_daily_lag14_zhengfawei$First_Infected), "%d %b, %Y")

regional_results_daily_lag14_zhengfawei <- regional_results_daily_lag14_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

regional_mobilization_lag14_zhengfawei <- ggplot(regional_results_daily_lag14_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  
  scale_x_continuous(breaks = regional_results_daily_lag14_zhengfawei$Order, labels = regional_results_daily_lag14_zhengfawei$City) +
  labs(title = "",
       x = "",
       y = "Regional Mobilization, Lagged 14 Days") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        # Hide middle-panel axes to declutter
        axis.text.y = element_blank(),
        axis.ticks.y = element_blank(),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

print(regional_mobilization_lag14_zhengfawei)

# Plot regional mobilization, daily, lag 21 days
regional_results_daily_lag21_zhengfawei <- regional_results_daily_lag21_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

regional_results_daily_lag21_zhengfawei$First_Infected <- format(as.Date(regional_results_daily_lag21_zhengfawei$First_Infected), "%d %b, %Y")

regional_results_daily_lag21_zhengfawei <- regional_results_daily_lag21_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

regional_mobilization_lag21_zhengfawei <- ggplot(regional_results_daily_lag21_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  
  scale_x_continuous(breaks = regional_results_daily_lag21_zhengfawei$Order, labels = regional_results_daily_lag21_zhengfawei$City, position = "top") +
  labs(title = "",
       x = "",
       y = "Regional Mobilization, Lagged 21 Days") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        axis.text.y = element_text(size = 16),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

regional_mobilization_lag21_zhengfawei


title_textA11_2 <- "FIGURE A11_2. Lagged Mobilization Effects at the Regional Level (Robustness Checks)"
title_grobA11_2 <- grobTree(
  textGrob(title_textA11_2, 
           gp = gpar(fontsize = 24, fontface = "bold"), y = -0.6)
)

# Figure notes (bottom)
note_textA11_2 <- "\nThis plot presents the lagged mobilization effects at the regional level. The size of each point indicates the scale of the local outbreak. Solid \npoints represent outbreaks that spread to other cities, while hollow points represent those that did not. Points colored in black indicate the \nstatistical significance of each mobilization effect. The left side indicates the date the first infection was detected, and the right side indicates the \ncity where the outbreak occurred."
note_grobA11_2 <- textGrob(note_textA11_2, 
                           gp = gpar(fontsize = 16), 
                           x = 0.1, hjust = 0.0)


figureA11_2 <- arrangeGrob(
  regional_mobilization_lag7_zhengfawei,
  regional_mobilization_lag14_zhengfawei,
  regional_mobilization_lag21_zhengfawei,
  ncol = 3,
  widths = c(1.08, 0.84, 1.08) 
)

ggsave(filename = "plots/figureA11_2.png", plot = figureA11_2, width = 16.5, height = 9.2, units = "in", dpi = 200)

# National mobilization with lags
# Lag = 7 days
process_data_national_lag7_zhengfawei <- function(test_list, shengjizhengfa) {
  all_results <- vector("list", nrow(test_list))
  
  for (i in seq_len(nrow(test_list))) {
    First_Infected     <- as.Date(test_list$First_Infected[i])
    City               <- test_list$City[i]
    Scale              <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date   <- First_Infected %m+% months(12)
    
    nl_raw   <- test_list$National_level[i]
    national <- if (is.na(nl_raw) || nl_raw == "") character(0) else {
      trimws(unlist(strsplit(nl_raw, ",")))
    }
    
    rd_data <- shengjizhengfa |>
      mutate(date = as.Date(date)) |>   
      filter(!(province %in% national),
             date >= start_date, date <= end_date) |>
      group_by(date) |>
      summarise(
        total_count = sum(total_count, na.rm = TRUE),
        daily_grid  = sum(daily_grid,  na.rm = TRUE),
        .groups = "drop"
      ) |>
      complete(date = seq.Date(start_date, end_date, by = "day"),
               fill = list(total_count = 0L, daily_grid = 0L)) |>
      arrange(date) |>
      mutate(
        difference = as.integer(date - First_Infected),
        others     = pmax(total_count - daily_grid, 0L)
      )
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    if (is.na(daily_total_mean) || daily_total_mean == 0 || nrow(rd_data) == 0) {
      coef <- NA_real_; ci_l <- NA_real_; ci_u <- NA_real_
    } else {
      rd_control <- rdrobust(
        y = rd_data$daily_grid,
        x = rd_data$difference,
        covs = rd_data$others,
        p = 1,
        kernel = "triangular",
        bwselect = "mserd",
        c = 7
      )
      coef <- rd_control[["coef"]]["Conventional", "Coeff"]  / daily_total_mean
      ci_l <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
      ci_u <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    }
    
    all_results[[i]] <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = coef,
      CI_Lower = ci_l,
      CI_Upper = ci_u,
      stringsAsFactors = FALSE
    )
  }
  
  dplyr::bind_rows(all_results)
}


national_results_daily_lag7_zhengfawei <- process_data_national_lag7_zhengfawei(test_list, shengjizhengfa)

print(national_results_daily_lag7_zhengfawei)

save(national_results_daily_lag7_zhengfawei, file = "generated_data/dynamic_RDD_results/national_results_daily_lag7_zhengfawei.RData")

# Lag = 14 days
process_data_national_lag14_zhengfawei <- function(test_list, shengjizhengfa) {
  all_results <- vector("list", nrow(test_list))
  
  for (i in seq_len(nrow(test_list))) {
    First_Infected     <- as.Date(test_list$First_Infected[i])
    City               <- test_list$City[i]
    Scale              <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date   <- First_Infected %m+% months(12)
    
    nl_raw   <- test_list$National_level[i]
    national <- if (is.na(nl_raw) || nl_raw == "") character(0) else {
      trimws(unlist(strsplit(nl_raw, ",")))
    }
    
    rd_data <- shengjizhengfa |>
      mutate(date = as.Date(date)) |>   
      filter(!(province %in% national),
             date >= start_date, date <= end_date) |>
      group_by(date) |>
      summarise(
        total_count = sum(total_count, na.rm = TRUE),
        daily_grid  = sum(daily_grid,  na.rm = TRUE),
        .groups = "drop"
      ) |>
      complete(date = seq.Date(start_date, end_date, by = "day"),
               fill = list(total_count = 0L, daily_grid = 0L)) |>
      arrange(date) |>
      mutate(
        difference = as.integer(date - First_Infected),
        others     = pmax(total_count - daily_grid, 0L)
      )
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    if (is.na(daily_total_mean) || daily_total_mean == 0 || nrow(rd_data) == 0) {
      coef <- NA_real_; ci_l <- NA_real_; ci_u <- NA_real_
    } else {
      rd_control <- rdrobust(
        y = rd_data$daily_grid,
        x = rd_data$difference,
        covs = rd_data$others,
        p = 1,
        kernel = "triangular",
        bwselect = "mserd",
        c = 14
      )
      coef <- rd_control[["coef"]]["Conventional", "Coeff"]  / daily_total_mean
      ci_l <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
      ci_u <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    }
    
    all_results[[i]] <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = coef,
      CI_Lower = ci_l,
      CI_Upper = ci_u,
      stringsAsFactors = FALSE
    )
  }
  
  dplyr::bind_rows(all_results)
}

national_results_daily_lag14_zhengfawei <- process_data_national_lag14_zhengfawei(test_list, shengjizhengfa)

print(national_results_daily_lag14_zhengfawei)

save(national_results_daily_lag14_zhengfawei, file = "generated_data/dynamic_RDD_results/national_results_daily_lag14_zhengfawei.RData")

# Lag = 21 days
process_data_national_lag21_zhengfawei <- function(test_list, shengjizhengfa) {
  all_results <- vector("list", nrow(test_list))
  
  for (i in seq_len(nrow(test_list))) {
    First_Infected     <- as.Date(test_list$First_Infected[i])
    City               <- test_list$City[i]
    Scale              <- test_list$Scale[i]
    Regional_spreading <- test_list$Regional_spreading[i]
    National_spreading <- test_list$National_spreading[i]
    
    start_date <- First_Infected %m-% months(12)
    end_date   <- First_Infected %m+% months(12)
    
    nl_raw   <- test_list$National_level[i]
    national <- if (is.na(nl_raw) || nl_raw == "") character(0) else {
      trimws(unlist(strsplit(nl_raw, ",")))
    }
    
    rd_data <- shengjizhengfa |>
      mutate(date = as.Date(date)) |>   
      filter(!(province %in% national),
             date >= start_date, date <= end_date) |>
      group_by(date) |>
      summarise(
        total_count = sum(total_count, na.rm = TRUE),
        daily_grid  = sum(daily_grid,  na.rm = TRUE),
        .groups = "drop"
      ) |>
      complete(date = seq.Date(start_date, end_date, by = "day"),
               fill = list(total_count = 0L, daily_grid = 0L)) |>
      arrange(date) |>
      mutate(
        difference = as.integer(date - First_Infected),
        others     = pmax(total_count - daily_grid, 0L)
      )
    
    daily_total_mean <- mean(rd_data$total_count, na.rm = TRUE)
    
    if (is.na(daily_total_mean) || daily_total_mean == 0 || nrow(rd_data) == 0) {
      coef <- NA_real_; ci_l <- NA_real_; ci_u <- NA_real_
    } else {
      rd_control <- rdrobust(
        y = rd_data$daily_grid,
        x = rd_data$difference,
        covs = rd_data$others,
        p = 1,
        kernel = "triangular",
        bwselect = "mserd",
        c = 21
      )
      coef <- rd_control[["coef"]]["Conventional", "Coeff"]  / daily_total_mean
      ci_l <- rd_control[["ci"]]["Conventional", "CI Lower"] / daily_total_mean
      ci_u <- rd_control[["ci"]]["Conventional", "CI Upper"] / daily_total_mean
    }
    
    all_results[[i]] <- data.frame(
      City = City,
      Scale = Scale,
      Regional_spreading = Regional_spreading,
      National_spreading = National_spreading,
      First_Infected = First_Infected,
      Coefficient = coef,
      CI_Lower = ci_l,
      CI_Upper = ci_u,
      stringsAsFactors = FALSE
    )
  }
  
  dplyr::bind_rows(all_results)
}

national_results_daily_lag21_zhengfawei <- process_data_national_lag21_zhengfawei(test_list, shengjizhengfa)

print(national_results_daily_lag21_zhengfawei)

save(national_results_daily_lag21_zhengfawei, file = "generated_data/dynamic_RDD_results/national_results_daily_lag21_zhengfawei.RData")

# Plot national mobilization, daily, lag 7 days
national_results_daily_lag7_zhengfawei <- national_results_daily_lag7_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

national_results_daily_lag7_zhengfawei$First_Infected <- format(as.Date(national_results_daily_lag7_zhengfawei$First_Infected), "%d %b, %Y")

national_results_daily_lag7_zhengfawei <- national_results_daily_lag7_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

national_mobilization_lag7_zhengfawei <- ggplot(national_results_daily_lag7_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  
  scale_x_continuous(breaks = national_results_daily_lag7_zhengfawei$Order, labels = national_results_daily_lag7_zhengfawei$First_Infected) +
  labs(title = "",
       x = "",
       y = "National Mobilization, Lagged 7 Days") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        axis.text.y = element_text(size = 16),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

print(national_mobilization_lag7_zhengfawei)

# Plot national mobilization, daily, lagged 14 days
national_results_daily_lag14_zhengfawei <- national_results_daily_lag14_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

national_results_daily_lag14_zhengfawei$First_Infected <- format(as.Date(national_results_daily_lag14_zhengfawei$First_Infected), "%d %b, %Y")

national_results_daily_lag14_zhengfawei <- national_results_daily_lag14_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

national_mobilization_lag14_zhengfawei <- ggplot(national_results_daily_lag14_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  
  scale_x_continuous(breaks = national_results_daily_lag14_zhengfawei$Order, labels = national_results_daily_lag14_zhengfawei$City) +
  labs(title = "",
       x = "",
       y = "National Mobilization, Lagged 14 Days") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        # Hide middle-panel axes to declutter
        axis.text.y = element_blank(),
        axis.ticks.y = element_blank(),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

print(national_mobilization_lag14_zhengfawei)

# Plot national mobilization, daily, lag 21 days
national_results_daily_lag21_zhengfawei <- national_results_daily_lag21_zhengfawei |>
  arrange(desc(First_Infected)) |>
  mutate(Order = row_number())

national_results_daily_lag21_zhengfawei$First_Infected <- format(as.Date(national_results_daily_lag21_zhengfawei$First_Infected), "%d %b, %Y")

national_results_daily_lag21_zhengfawei <- national_results_daily_lag21_zhengfawei |>
  mutate(Significance = ifelse(CI_Lower > 0 & CI_Upper > 0, 1, 0))

national_mobilization_lag21_zhengfawei <- ggplot(national_results_daily_lag21_zhengfawei, aes(x = Order, y = Coefficient)) +
  geom_point(aes(size = Scale * 2, color = as.factor(Significance), 
                 fill = as.factor(Significance), shape = as.factor(Regional_spreading)), 
             stroke = 2) +
  geom_errorbar(aes(ymin = CI_Lower, ymax = CI_Upper, color = as.factor(Significance)), 
                width = 0.4, linewidth = 0.8, linetype = "solid") +
  geom_hline(yintercept = 0, color = "#0033FF", linewidth = 1, linetype = "dashed") +
  scale_size_continuous(range = c(2, 10)) +
  scale_color_manual(values = c("0" = "lightgray", "1" = "black")) +
  scale_fill_manual(values = c("0" = "lightgray", "1" = "black")) + 
  scale_shape_manual(values = c("0" = 1, "1" = 16)) +  # Empty vs solid points by spreading
  scale_x_continuous(breaks = national_results_daily_lag21_zhengfawei$Order, labels = national_results_daily_lag21_zhengfawei$City, position = "top") +
  labs(title = "",
       x = "",
       y = "National Mobilization, Lagged 21 Days") +
  theme_bw() +
  theme(legend.position = "none",
        text = element_text(size = 18),
        axis.text.y = element_text(size = 16),
        axis.text.x = element_blank(),  
        axis.ticks.x = element_blank(),  
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill = NA, linewidth = 1)) +
  coord_flip()

national_mobilization_lag21_zhengfawei

title_textA11_3 <- "FIGURE A11_3. Lagged Mobilization Effect at the National Level (Robustness Checks)"
title_grobA11_3 <- grobTree(
  textGrob(title_textA11_3, 
           gp = gpar(fontsize = 24, fontface = "bold"), y = -0.6)
)

# Figure notes (bottom)
note_textA11_3 <- "\nThis plot presents the lagged mobilization effects at the national level. The size of each point indicates the scale of the local outbreak. Solid \npoints represent outbreaks that spread to other cities, while hollow points represent those that did not. Points colored in black indicate the \nstatistical significance of each mobilization effect. The left side indicates the date the first infection was detected, and the right side indicates the \ncity where the outbreak occurred."
note_grobA11_3 <- textGrob(note_textA11_3, 
                           gp = gpar(fontsize = 16), 
                           x = 0.1, hjust = 0.0)


figureA11_3 <- arrangeGrob(
  national_mobilization_lag7_zhengfawei,
  national_mobilization_lag14_zhengfawei,
  national_mobilization_lag21_zhengfawei,
  ncol = 3,
  widths = c(1.08, 0.84, 1.08) 
)

ggsave(filename = "plots/FIGUREA11_3.png", plot = figureA11_3, width = 16.5, height = 9.2, units = "in", dpi = 200)

# Combine Figure A11_2 and Figure A11_3 into a single tall panel
figure_combinedA11_2_3 <- grid.arrange(figureA11_2, figureA11_3, ncol = 1)
ggsave(filename = "plots/FIGURE_combinedA11_2_3.png", plot = figure_combinedA11_2_3, width = 16.5, height = 18.4, units = "in", dpi = 200)
